#!/usr/bin/env python3

import rospy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError
import cv2

import numpy as np

from sympy import Point, Line, Point2D

from yolo_detector import YoloDetector
from pointer_reader import PointerDialReader
from pointer_value_convertor import PointerValueConvertor


class ImageViewer:
    def __init__(self):
        # 初始化ros node
        rospy.init_node('leak_detection', anonymous=True)

        # 创建bridge对象以便于转化msg和opencv images
        self.bridge = CvBridge()

        # 订阅主题 /usb_cam/image_raw
        self.image_sub = rospy.Subscriber("/usb_cam/image_raw",Image,self.callback)

        # 黑色像素计数器
        self.black_pixel_count = 0

        # 设置关闭窗口事件监听
        cv2.startWindowThread()

        # 仪表盘参数
        self.dashboard_params = {
            'min_angle': 225,  # 仪表最小量程时，指针角度（指针竖直向上是为0，顺时针旋转）
            'max_angle': 135,  # 仪表最大量程时，指针角度（指针竖直向上是为0，顺时针旋转）
            'min_value': 0,  # 仪表最小量程
            'max_value': 2.5,  # 仪表最大量程
        }

        # 初始化 YOLO 检测器
        self.yolo_detector = YoloDetector()

        # 初始化表盘读取器
        self.dial_reader = PointerDialReader()

        # 初始化指针数据转换
        self.convertor = PointerValueConvertor(
            self.dashboard_params['min_angle'],
            self.dashboard_params['max_angle'],
            self.dashboard_params['min_value'],
            self.dashboard_params['max_value']
        )

        # YOLO 预读取帧数设置
        self.loaded_time = 0
        self.pre_load_limit = self.yolo_detector.get_worker_number() + 1

    def callback(self,data):
        try:
            # 使用桥将ros msg转化为opencv image
            frame = self.bridge.imgmsg_to_cv2(data,"bgr8")

            # 初始化异步所需要的帧
            if self.loaded_time < self.pre_load_limit:
                self.yolo_detector.fill_pool(frame)
                self.loaded_time += 1
                return

            # 检测表盘
            detections = self.yolo_detector.process_frame(frame)

            # 跳过没有检测到表盘的数据
            if len(detections) < 1:
                return

            # 解包检测后的数据
            name, score, watcher_center, boundary = detections[0]
            x1, y1, x2, y2 = boundary

            # 跳过识别有问题的帧
            if (x2 - x1) < 0 or (y2 - y1) < 0:
                return

            # 截取表盘区域
            watcher_roi = frame[y1:y2, x1:x2]

            # 检测表盘中心和半径
            center, radius = self.dial_reader.find_dial_center(watcher_roi)
            # print(f"center: {center}, radius: {radius}")

            if center is not None:
                cv2.circle(watcher_roi, center, 4, (0, 0, 255), 2)
                cv2.circle(watcher_roi, center, radius, (0, 0, 255), 2)

            # 检测指针
            pointer = self.dial_reader.detect_pointer(watcher_roi, center, radius)
            # print(f"pointer: {pointer}")

            if pointer is not None:
                cv2.line(watcher_roi, (pointer[0], pointer[1]), (pointer[2], pointer[3]), (255, 0, 0), 2)

            # 跳过未检测到表盘中心和指针的帧
            if center is None or pointer is None:
                return

            # 计算指针角度
            angle = self.dial_reader.calculate_angle(pointer, center)
            print(f"angle: {angle}")

            # 将指针角度转换为表盘上的读数
            value = self.convertor.angle_to_value(angle)
            print(f"value: {value}")

            # 显示读数
            if 0 < value < 0.85:
                cv2.putText(
                    watcher_roi,
                    text=f"{value}",
                    org=(10, 30),
                    fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=1.0,
                    color=(0, 255, 255),
                    thickness=2
                )

            elif 0.85 <= value < 1.65:
                cv2.putText(
                    watcher_roi,
                    text=f"{value}",
                    org=(10, 30),
                    fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=1.0,
                    color=(0, 255, 0),
                    thickness=2
                )

            elif 1.65 <= value <= 2.5:
                cv2.putText(
                    watcher_roi,
                    text=f"{value}",
                    org=(10, 30),
                    fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=1.0,
                    color=(0, 0, 255),
                    thickness=2
                )

            else:
                cv2.putText(
                    watcher_roi,
                    text="unknown",
                    org=(10, 30),
                    fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=1.0,
                    color=(255, 0, 0),
                    thickness=2
                )

            cv2.namedWindow("watcher_roi", cv2.WINDOW_NORMAL)
            cv2.imshow("watcher_roi", watcher_roi)

        except CvBridgeError as e:
            print(e)


if __name__=='__main__':
    viewer=ImageViewer()

    try:
        rospy.spin()
    finally:
        cv2.destroyAllWindows()